package net.lab0.fractal.fractals.mandelbrot;

import static com.jogamp.opencl.CLMemory.Mem.READ_ONLY;
import static com.jogamp.opencl.CLMemory.Mem.WRITE_ONLY;
import static java.lang.Math.min;
import static java.lang.System.nanoTime;
import static java.lang.System.out;

import java.io.IOException;
import java.nio.FloatBuffer;
import java.nio.IntBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;

import com.jogamp.opencl.CLBuffer;
import com.jogamp.opencl.CLCommandQueue;
import com.jogamp.opencl.CLContext;
import com.jogamp.opencl.CLDevice;
import com.jogamp.opencl.CLKernel;
import com.jogamp.opencl.CLProgram;

public class OpenCLMandelbrotComputing
{
    public OpenCLMandelbrotComputing()
    {
        // set up (uses default CLPlatform and creates context for all devices)
        CLContext context = CLContext.create();
        out.println("created " + context);
        
        // always make sure to release the context under all circumstances
        // not needed for this particular sample but recommented
        try
        {
            // select fastest device
            CLDevice device = context.getMaxFlopsDevice();
            out.println("using " + device);
            
            // create command queue on device.
            CLCommandQueue queue = device.createCommandQueue();
            
            int maxIter = 1 << 16;
            int elementCount = 10000; // Length of arrays to process
            int localWorkSize = min(device.getMaxWorkGroupSize(), 256); // Local work size dimensions
            int globalWorkSize = roundUp(localWorkSize, elementCount); // rounded up to the nearest multiple of the localWorkSize
            
            // load sources, create and build program
            CLProgram program = context.createProgram(OpenCLMandelbrotComputing.class.getResourceAsStream("programm.cl ")).build();
            
            long time = nanoTime();
            // A, B are input buffers, C is for the result
            CLBuffer<FloatBuffer> clBufferA = context.createFloatBuffer(globalWorkSize, READ_ONLY);
            CLBuffer<FloatBuffer> clBufferB = context.createFloatBuffer(globalWorkSize, READ_ONLY);
            CLBuffer<IntBuffer> clBufferC = context.createIntBuffer(globalWorkSize, WRITE_ONLY);
            
            out.println("used device memory: " + (clBufferA.getCLSize() + clBufferB.getCLSize() + clBufferC.getCLSize()) / 1024 / 1024 + "MiB");
            
            // fill input buffers with random numbers
            // (just to have test data; seed is fixed -> results will not change between runs).
            fillBuffer(clBufferA.getBuffer(), 12345);
            fillBuffer(clBufferB.getBuffer(), 67890);
            
            // get a reference to the kernel function with the name 'VectorAdd'
            // and map the buffers to its input parameters.
            CLKernel kernel = program.createCLKernel("f");
            kernel.putArgs(clBufferA, clBufferB, clBufferC).putArg(maxIter);
            
            // asynchronous write of data to GPU device,
            // followed by blocking read to get the computed results back.
            time = nanoTime() - time;
            out.println("init: " + (time / 1000000) + "ms");
            time = nanoTime();
            queue.putWriteBuffer(clBufferA, false).putWriteBuffer(clBufferB, false).put1DRangeKernel(kernel, 0, globalWorkSize, localWorkSize)
                    .putReadBuffer(clBufferC, true);
            time = nanoTime() - time;
            
            List<Integer> gpuResults = new ArrayList<Integer>(elementCount);
            // print first few elements of the resulting buffer to the console.
            out.println("results snapshot: ");
            for (int i = 0; i < 10; i++)
            {
                int val = clBufferC.getBuffer().get();
                out.print(val + ", ");
                gpuResults.add(val);
            }
            out.println("...; " + clBufferC.getBuffer().remaining() + " more");
            
            int remainder = clBufferC.getBuffer().remaining();
            for (int i = 0; i < remainder; i++)
            {
                int val = clBufferC.getBuffer().get();
                gpuResults.add(val);
            }
            
            out.println("openCL computation took: " + (time / 1000000) + "ms");
            
            int remaining = clBufferC.getBuffer().remaining();
            int max = 0;
            for (int i = 0; i < remaining; i++)
            {
                int value = clBufferC.getBuffer().get();
                if (value > max)
                {
                    max = value;
                }
            }
            System.out.println("max: " + max);
            
            time = nanoTime();
            cpuMandelbrot(clBufferA, clBufferB, clBufferC, maxIter);
            time = nanoTime() - time;
            
            List<Integer> cpuResults = new ArrayList<Integer>(elementCount);
            clBufferC.getBuffer().rewind();
            // print first few elements of the resulting buffer to the console.
            out.println("results snapshot: ");
            for (int i = 0; i < 10; i++)
            {
                int val = clBufferC.getBuffer().get();
                out.print(val + ", ");
                cpuResults.add(val);
            }
            out.println("...; " + clBufferC.getBuffer().remaining() + " more");
            
            out.println("CPU computation took: " + (time / 1000000) + "ms");
            
            remaining = clBufferC.getBuffer().remaining();
            max = 0;
            for (int i = 0; i < remaining; i++)
            {
                int value = clBufferC.getBuffer().get();
                cpuResults.add(value);
                if (value > max)
                {
                    max = value;
                }
            }
            System.out.println("max: " + max);
            
            for (int i = 0; i < elementCount; ++i)
            {
                if (cpuResults.get(i).intValue() != gpuResults.get(i).intValue())
                {
                    System.out.println("Diff@" + i + " : " + cpuResults.get(i) + "/" + gpuResults.get(i) + "->" + (cpuResults.get(i) - gpuResults.get(i)));
                }
            }
        }
        catch (IOException e)
        {
            e.printStackTrace();
        }
        finally
        {
            // cleanup all resources associated with this context.
            context.release();
        }
    }
    
    private void cpuMandelbrot(CLBuffer<FloatBuffer> clBufferA, CLBuffer<FloatBuffer> clBufferB, CLBuffer<IntBuffer> clBufferC, int max)
    {
        clBufferA.getBuffer().rewind();
        clBufferB.getBuffer().rewind();
        clBufferC.getBuffer().rewind();
        int size = clBufferA.getBuffer().remaining();
        for (int i = 0; i < size; ++i)
        {
            float realPos = clBufferA.getBuffer().get();
            float imagPos = clBufferB.getBuffer().get();
            float real1 = realPos;
            float img1 = imagPos;
            float real2, img2;
            float realSquared = real1 * real1;
            float imagSquared = img1 * img1;
            
            int iter = 0;
            while ((iter < max) && ((realSquared + imagSquared) < 4))
            {
                real2 = real1 * real1 - img1 * img1 + realPos;
                img2 = 2 * real1 * img1 + imagPos;
                
                real1 = real2 * real2 - img2 * img2 + realPos;
                img1 = 2 * real2 * img2 + imagPos;
                
                realSquared = real2 * real2;
                imagSquared = img2 * img2;
                
                iter += 2;
            }
            clBufferC.getBuffer().put(iter);
        }
    }
    
    private void fillBuffer(FloatBuffer buffer, int seed)
    {
        Random rnd = new Random(seed);
        while (buffer.remaining() != 0)
            buffer.put(rnd.nextFloat() * 2.0f - 1.0f);
        buffer.rewind();
    }
    
    private int roundUp(int groupSize, int globalSize)
    {
        int r = globalSize % groupSize;
        if (r == 0)
        {
            return globalSize;
        }
        else
        {
            return globalSize + groupSize - r;
        }
    }
    
    public static void main(String[] args)
    {
        new OpenCLMandelbrotComputing();
    }
}
